'''
这里分别构建主要的深度学习的问题
'''
import matplotlib.pyplot as plt
import numpy as np

from skimage.data import astronaut
from skimage.color import rgb2gray
from skimage import filters as filters
from skimage.segmentation import felzenszwalb, slic, quickshift, watershed
from skimage.segmentation import mark_boundaries
from skimage.util import img_as_float
import torch
from models.get_models import get_models
import argparse
from utils.utils import get_optimizer,get_scheduler,save_ckpt,save_ckpt_bestmiou,load_ckpt
from utils.loginfomation import loginfomation,WriterAccurary,WriterSummary
from utils.loss import *

from datasetloader.get_datasetLoader import *
import time
from utils.metrics import runningScore
from utils.loss import get_lossfunction
from torchvision.utils import make_grid
import os
import torch.nn.functional as F
import sys
import threading
from time import sleep

from scipy import io

from config.config_fcn import fcn_model_config_RGB123
from config.config_unet import unet_model_config_RGB123
from config.config_segnet import segnet_model_config_RGB123


def save_model(running_metrics):
    scores_val, class_iou_val = running_metrics.get_scores()
    for k, v in scores_val.items():
        print(k+': %f' % v)

    for k, v in class_iou_val.items():
        print(str(k)+': %f' % v)    
    running_metrics.reset()
    return scores_val,class_iou_val,running_metrics



def predict(config_param,modelName,cls_num,channels_num=3,redmodelpath=""):

    E_train_loader, test_loader = get_dataset(config_param,dataSize='E512') 
    model=get_models(modelName,num_class=cls_num,channels_num=3)
    startepoch,model,best_iou=load_ckpt(model, None, redmodelpath,'cuda')
    model.cuda()
    model.eval()        
    model.eval()
    strlines = []
    epoch=0
    running_metrics = runningScore(cls_num)
    with torch.no_grad():
        for i,sample in enumerate(test_loader):
                end = time.time()
                image,label,source_img,seg=sample['img'],  sample['label'],sample["source_image"],sample['seg']
                images = image.cuda()
                outputs = model(image.cuda())
                _, _, image_h, image_w = label.shape
                pred = F.softmax(outputs).cpu()  # 因为原始模型中结尾 没有softmax
                # pred = torch.argmax(pred, dim=1).cpu()
                gt = label.data.cpu()
                # 统计损失值
                running_metrics.update(gt.numpy(), pred.numpy() > 0.5)

        '''模型度量阶段'''
        scores_val,class_iou_val,running_metrics=save_model(running_metrics)

def superPixel(img_path):
    img=np.load(img_path)
    segments_fz = felzenszwalb(img, scale=100, sigma=0.5, min_size=50)
    segments_slic = slic(img, n_segments=250, compactness=10, sigma=1)
    segments_quick = quickshift(img, kernel_size=3, max_dist=6, ratio=0.5)
    gradient = filters.sobel(rgb2gray(img))
    segments_watershed = watershed(gradient, markers=250, compactness=0.001)
    return segments_fz,segments_slic,segments_quick,segments_watershed

def arigmation(fcn_pred,superpixel,seg_pred,img,a,b):

    mean_pred=(a*fcn_pred+b*seg_pred)/2
    segments_quick = quickshift(img, kernel_size=3, max_dist=6, ratio=0.5)
    superpixel_mask_all=segments_quick*0
    ids=np.unique(segments_quick)
    for i in ids:
        spp_mask=segments_quick==i
        mean_pred_mask=spp_mask*mean_pred
        mean_preds=np.sum(np.sum(mean_pred_mask))/np.sum(np.sum(spp_mask))
        if mean_preds>0.5:
            superpixel_mask_all=superpixel_mask_all+spp_mask
    return superpixel_mask_all


# 问题主流程

def mainprocess():
    imgpath='/media/gis/databackup/ayc/modellist/dataset/nanchang/E512/RGB123/test/img/1818.npy'
    labelpath='/media/gis/databackup/ayc/modellist/dataset/nanchang/E512/RGB123/test/label/1818.npy'
    fcn_model_path='/media/gis/databackup/ayc/model_hastrained_log/RGB123_fcn_result/checkpoint/Bestest_mIou_ckpt_fcn_dataName_RGB123.pth'
    seg_model_path='/media/gis/databackup/ayc/model_hastrained_log/RGB123_segnet_result/checkpoint/Bestest_mIou_ckpt_segnet_dataName_RGB123.pth'
    unet_model_path='/media/gis/databackup/ayc/model_hastrained_log/RGB123_unet_result/checkpoint/Bestest_mIou_ckpt_unet_dataName_RGB123.pth'
    # 模型预测阶段
    predict(fcn_model_config_RGB123,"fcn",2,channels_num=3,redmodelpath=fcn_model_path)
    predict(unet_model_config_RGB123,"unet",2,channels_num=3,redmodelpath=unet_model_path)

if __name__=="__main__":
    mainprocess()


